import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt


class Net(torch.nn.Module):

    def __init__(self, n_features, n_hidden, n_output):  # 搭建神经网络层的信息
        super(Net, self).__init__()
        self.hidden = torch.nn.Linear(n_features, n_hidden)  # 隐藏层
        self.predict = torch.nn.Linear(n_hidden, n_output)  # 输出层

    def forward(self, x):  # 神经网络前向传递，也即搭建神经网络
        x = F.relu(self.hidden(x))
        x = self.predict(x)
        return x


if __name__ == '__main__':

    x = torch.unsqueeze(torch.linspace(-1, 1, 100), dim=1)
    y = x.pow(2) + 0.2 * torch.rand(x.size())

    net = Net(1, 10, 1)

    optimizer = torch.optim.SGD(net.parameters(), lr=0.5)
    loss_func = torch.nn.MSELoss()

    plt.figure()
    plt.ion()  # plt一个实时打印的状态

    for t in range(200):
        prediction = net(x)
        loss = loss_func(prediction, y)

        optimizer.zero_grad()  # 将梯度归零
        loss.backward()
        optimizer.step()

        if t % 10 == 0:
            plt.cla()
            plt.scatter(x.numpy(), y.numpy())
            plt.plot(x.detach().numpy(), prediction.detach().numpy(), 'r-', lw=5)
            plt.text(0.5, 0, 'Loss=%.4f' % loss, fontdict={'size': 10, 'color': 'red'})
            plt.pause(0.1)

    plt.ioff()
    plt.show()
